{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 具体函数\n",
    "在 [AutoGraph and tf.functions](https://tensorflow.google.cn/guide/function) 中，你看到了如何使用 [tf.funtion](https://tensorflow.google.cn/api_docs/python/tf/function)。这篇指南将详细介绍：\n",
    "- tf.funtion 跟踪\n",
    "- tf.funtion 签名\n",
    "- 跟踪生成的具体函数\n",
    "    - 如何访问它们\n",
    "    - 如何使用它们\n",
    "这些细节只会变得重要：\n",
    "- 如果由于不希望跟踪tf.function而遇到性能问题。\n",
    "- 当需要精确控制 tf.function 生成的 tensorflow 图时。例如，使用 tf.Lite.Converter.from_concrete_functions 将模型导出到 [TensorFlow Lite](https://tensorflow.org/lite/)。\n",
    "\n",
    "## 出现背景\n",
    "在 tensorflow2 中，默认情况下 \"eager execution\" 处于启用状态。tensorflow 的 eager exectuion 是一个命令式编程环境，它可以立即评估操作，而无需构建图。它操作返回值而不是构造计算图以便以后运行。这里可查看 eager execution 详细信息 [detailed guide on eager execution](https://tensorflow.google.cn/guide/eager)。\n",
    "\n",
    "运行命令使开发和调试更具交互性，但它不能轻松地导出模型。\n",
    "\n",
    "tf.function API 可以将模型保存为图形。\n",
    "\n",
    "## 相关术语\n",
    "下面的术语会在文中被用到：\n",
    "- 签名：一组操作的输入和输出的描述。\n",
    "- 多态函数：可调用的 python，它在一个 AP I后面封装了几个具体的函数图。\n",
    "- 具体函数：带有单个签名的图。\n",
    "\n",
    "## 设置"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import traceback\n",
    "import textwrap\n",
    "\n",
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 创建一个 tf.funtion\n",
    "tf.function 的注释函数会生成包含这些操作的多态函数。所有未用 tf.function 注释的操作都将使用 eager execution 进行计算。下面的示例显示了 tf.function 用法的一个快速示例。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tf.function\n",
    "def square(x):\n",
    "  return x*x\n",
    "\n",
    "square(2).numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 记住：python 的装饰语法只是调用装饰对象作为输入装饰\n",
    "def pow(x,y):\n",
    "  return x ** y\n",
    "\n",
    "pow = tf.function(pow)\n",
    "\n",
    "pow(3,4).numpy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 将 tf.function 附加 tf.Module 上\n",
    "tf.function 可以选择性地存储为 tf.Module 对象的一部分。Module 类提供了跟踪变量、保存检查点和模型的特性。\n",
    "\n",
    "类如 keras.layers.Layer 和 keras.Model 是模块的子类。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Pow(tf.Module):\n",
    "  def __init__(self, exponent):\n",
    "    self.exponent = tf.Variable(exponent, dtype = tf.float32, name='Pow/exponent')\n",
    "\n",
    "  @tf.function\n",
    "  def __call__(self, x):\n",
    "    return x ** self.exponent\n",
    "\n",
    "pow = Pow(3)\n",
    "\n",
    "print(pow.variables)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pow(tf.constant(2.0)).numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pow.exponent.assign(4)\n",
    "pow(tf.constant(2.0)).numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.saved_model.save(pow, 'pow')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reloaded_pow = tf.saved_model.load('pow')\n",
    "reloaded_pow(tf.constant(3.0)).numpy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 将 tf.function 指定为属性\n",
    "如果你将 tf.Module 或 tf.function 指定为模型的属性，则它也将被序列化："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mod = tf.Module()\n",
    "mod.increment_by = tf.Variable(2.0)\n",
    "\n",
    "@tf.function\n",
    "def increment(x):\n",
    "  return x+mod.increment_by\n",
    "\n",
    "mod.inc = increment\n",
    "mod.inc(tf.constant(1.0)).numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mod.cube = Pow(3)\n",
    "mod.cube(tf.constant(2.0)).numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(mod.variables)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.saved_model.save(mod, 'mod')\n",
    "reloaded_mod = tf.saved_model.load('mod')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reloaded_mod.inc(4.0).numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reloaded_mod.cube(4.0).numpy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 与 tf.keras 互动\n",
    "\n",
    "像 [Keras.Model](https://tensorflow.google.cn/api_docs/python/tf/keras/Model) 和 [Keras.layers.Layer](https://tensorflow.google.cn/api_docs/python/tf/keras/layers/Layer) 这样的 Keras 类与 tf.function 和 tf.Module 完全兼容。\n",
    "\n",
    "比如，构建一个简单的模型："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "linear = tf.keras.Sequential([tf.keras.layers.Dense(units=1, input_shape=[1])])\n",
    "linear.compile(optimizer='adam', loss='mean_squared_error')\n",
    "linear.fit(x=[-1, 0, 1, 2, 3, 4], y=[-3, -1, 1, 3, 5, 7], epochs=50, verbose=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "linear(tf.constant([[1],[2]]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 检查它的变量\n",
    "linear.variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 将它添加到 tf.Module 上\n",
    "module = tf.Module()\n",
    "module = linear"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# tf.Model 也跟踪 tf.Variable\n",
    "module.variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# tf.Model 还可导出 keras.Model 的内容\n",
    "tf.saved_model.save(module,'module')\n",
    "\n",
    "reloaded = tf.saved_model.load('module')\n",
    "reloaded.linear([[1.0]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 跟踪\n",
    "从 tf.function返 回的对象是多态函数。它们将接受 python 对象或 tf.Tensors（具有任意形状）或tf.dtype 作为输入。\n",
    "\n",
    "在后台，tensorflow 生成表示计算的 [tf.Graph](https://tensorflow.google.cn/api_docs/python/tf/Graph)。这个图包装在一个 python 可调用函数中：一个具体的函数。每个具体函数只能处理一个输入签名。\n",
    "\n",
    "tf.function 在每次需要创建具体函数时跟踪 python 函数。查看何时跟踪函数的最简单方法是添加对 print 的调用："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tf.function\n",
    "def mul(a, b):\n",
    "  print('Tracing:\\n    {a}\\n    {b}\\n'.format(a=a, b=b))\n",
    "  return a*b"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 数据类型和形状\n",
    "如果使用两种不同类型的输入调用多态函数，它将为每种类型跟踪一次："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#  使用 int 进行跟踪\n",
    "mul(tf.constant(2), tf.constant(3)).numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 使用 float 进行跟踪\n",
    "mul(tf.constant(2.0), tf.constant(3.0)).numpy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "当你使用相同的输入类型再次调用它时，它将分派给现有函数，而不是跟踪："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 再次调用 int\n",
    "mul(tf.constant(10), tf.constant(10))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "更改输入的大小也会触发跟踪（设置 [tf.function(experimental_relax_shapes=True)](https://tensorflow.google.cn/api_docs/python/tf/function) 可能会降低它的值）："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 使用向量进行跟踪\n",
    "mul(tf.constant([1.0,3.0]), tf.constant(3.0)).numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 使用不同大小的向量进行跟踪\n",
    "mul(tf.constant([1.0,2.0,3.0, 4.0]), tf.constant(3.0))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 不可变的 python 对象\n",
    "如果你传入了一个不可变的 python对象（如 int、st r或 tuple）给 tf.function，它将对这些 python 对象的每个值执行跟踪。\n",
    "\n",
    "这对于控制 tf.Graph 中包含的内容非常有用（参见：[The Autograph Guide](https://tensorflow.google.cn/guide/function) 获得更多信息）。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tf.function\n",
    "def mul(a, b):\n",
    "  print('跟踪:\\n    {a}\\n    {b}\\n'.format(a=a, b=b))\n",
    "  return a*b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 当 a=3.0 进行跟踪\n",
    "mul(3.0, tf.constant(3.0)).numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 当 a=3.0  不再跟踪\n",
    "mul(3.0, tf.constant(3.0)).numpy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> 注意：通过传递唯一的 pytho n值很容易导致许多跟踪。这可能是一个严重的性能问题。通常都传递一个tf.Variable 值。\n",
    "\n",
    "此循环跟踪每个唯一 int 的函数："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tf.function\n",
    "def power(a,b):\n",
    "  print('跟踪 : a={}'.format(a))\n",
    "  return a**b\n",
    "\n",
    "p = tf.constant(2)\n",
    "for n in range(12):\n",
    "  power(n,p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 在第二次运行时，每个 int 都被跟踪过了，因此不再需要跟踪\n",
    "p = tf.constant(2)\n",
    "for n in range(12):\n",
    "  power(n,p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 为了避免过度的回溯，应确保传递 tf.Tensor 而不是 python 数字或字符串：\n",
    "p = tf.constant(2)\n",
    "for n in tf.range(12):\n",
    "  power(n,p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 要完全关闭跟踪，将签名传递给 tf.function decorator\n",
    "@tf.function(input_signature=(\n",
    "    tf.TensorSpec(shape=[], dtype=tf.float32),\n",
    "    tf.TensorSpec(shape=[], dtype=tf.float32),)\n",
    ")\n",
    "def power_with_sig(a,b):\n",
    "  print('跟踪\"power_with_sig\"')\n",
    "  return a**b\n",
    "\n",
    "power_with_sig(3.0, 3.0).numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "  power_with_sig(tf.constant([1.0,2.0,3.0]),tf.constant(3.0))\n",
    "  assert False\n",
    "except ValueError:\n",
    "  traceback.print_exc(limit=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 实例：Dropout\n",
    "通过对特定值的回溯，可以控制 tf.function 生成的代码。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Dropout(tf.Module):\n",
    "  def __init__(self, rate, name=None):\n",
    "    super(Dropout, self).__init__(name)\n",
    "    self.rate = tf.Variable(rate, dtype = tf.float32, trainable=False)\n",
    "\n",
    "  @tf.function\n",
    "  def __call__(self, x, training=True):\n",
    "    print(textwrap.dedent(\"\"\"\n",
    "                          跟踪 \"Dropout\":\n",
    "                              训练 = {}\n",
    "                              x = {}\n",
    "                              名字 = {:s}\n",
    "                          \"\"\".format(training, x, self.name)))\n",
    "    if training:\n",
    "      print('    - 训练分支\\n')\n",
    "      mask = tf.random.uniform(x.shape) > self.rate\n",
    "      return x * tf.cast(mask, tf.float32)/self.rate\n",
    "    else:\n",
    "      print('    - 测试分支\\n')\n",
    "      return x\n",
    "    \n",
    "# 创建这个简单的 Dropout 层的一个实例\n",
    "dropout = Dropout(0.5)\n",
    "\n",
    "# 第一次使用 training=True 作为输入调用它时，它会跟踪训练分支\n",
    "dropout(tf.range(10, dtype=tf.float32), training=True).numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 第二次调用，它不需要再回溯分支了\n",
    "dropout(tf.range(10, dtype=tf.float32), training=True).numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 传递 training=False 将在第一次运行时触发跟踪，因为这是一个不同的值\n",
    "dropout(tf.range(10, dtype=tf.float32), training=False).numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dropout(tf.range(10, dtype=tf.float32), training=False).numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 如果你传递了一个 bool 张量，它将使用 TensorFlow autograph 将 If 重写为 tf.condm 并跟踪两个分支\n",
    "dropout(tf.range(10, dtype=tf.float32), training=tf.constant(False)).numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 这可以用于在单个具体函数中捕获控制流。\n",
    "dropout(tf.range(10, dtype=tf.float32), training=tf.constant(True)).numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dropout(tf.range(10, dtype=tf.float32), training=tf.constant(False)).numpy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 其它的 python 对象\n",
    "由于生成的 tf.Graphs 不能包含复杂的 python 对象，因此可以通过跟踪和变量捕获来包含这些对象。\n",
    "\n",
    "tf.function 为每个实例运行单独的跟踪。因此，每个跟踪都包含自己的变量，并且可以基于实例设置其行为。\n",
    "\n",
    "最常见的用法是关于模块、层或模块的方法："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dropout_a = Dropout(0.5, name='dropout_a')\n",
    "\n",
    "print(dropout_a(tf.range(10, dtype=tf.float32), True).numpy())\n",
    "print(dropout_a(tf.range(10, dtype=tf.float32), True).numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dropout_b = Dropout(0.5, name='dropout_b')\n",
    "\n",
    "print(dropout_b(tf.range(10, dtype=tf.float32), True).numpy())\n",
    "print(dropout_b(tf.range(10, dtype=tf.float32), True).numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 但是在一个独立的t f.function 上的行为是相同的。\n",
    "@tf.function\n",
    "def run(callable, x):\n",
    "  print('跟踪 \"run\":\\n    可调用 = {}\\n    x = {}\\n'.format(callable, x))\n",
    "  return callable(x)\n",
    "\n",
    "def plus_1(x):\n",
    "  return x+1\n",
    "\n",
    "print(run(plus_1, tf.constant(2.0)).numpy())\n",
    "print(run(plus_1, tf.constant(5.0)).numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 跟踪一个 tf.function 可以在另一个 tf.function 中触发跟踪\n",
    "print(run(dropout, tf.range(10.0)).numpy())\n",
    "print(run(dropout, tf.range(10.0)).numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 弱引用\n",
    "> 注意：每个跟踪只保留对 tf.Variable 的弱引用。如果变量没有被另一个引用保持活动状态，则跟踪可能变得不可用。\n",
    "\n",
    "例如，这里有一个 tf.function，它引用来自闭包的 var："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tf.function\n",
    "def plus_var(x):\n",
    "  print('跟踪 \"plus_var\":\\n    x = {}\\n    var = {}\\n\\n'.format(x, var.name))\n",
    "  return x + var"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 使用一个变量跟踪函数\n",
    "var = tf.Variable(1, name=\"IntVar\")\n",
    "plus_var(tf.constant([1,2])).numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 使用另一个变量\n",
    "var = tf.Variable(2.0, name=\"FloatVar\")\n",
    "plus_var(tf.constant([2.0, 10.0])).numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 这很有效，但是因为你不再有对 \"IntVar\" 的引用，所以第一个跟踪被破坏\n",
    "try:\n",
    "  plus_var(tf.constant([1,2])).numpy()\n",
    "  assert False\n",
    "except tf.errors.FailedPreconditionError:\n",
    "  traceback.print_exc(limit=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 访问具体函数\n",
    "在上一节中，你看到了触发多态 tf.function 的跟踪的条件。每个跟踪生成一个新的具体函数。\n",
    "\n",
    "当你将 tf.Module 保存为 tf.saved_model 时，将导出定义 tf.graph 的具体函数。你不保存 tf.function 而是保存通过跟踪创建的具体函数。\n",
    "\n",
    "要从多态 tf.function 获取具体的函数，需要定义签名。或者：\n",
    "- 将 input_signature 传递给 tf.function，并调用 get_concrete_function() 方法。\n",
    "- 传递 tf.TensorSpecs 的列表给 get_concrete_function：tf.TensorSpec(shape=[1], dtype=tf.float32)。\n",
    "- 传递一个正确形状和类型的示例张量以获得具体的函数：tf.constant(1., shape=[1]).。\n",
    "\n",
    "下面的示例演示如何定义 tf.function 的 input_signature 参数。\n",
    "\n",
    "**使用 input_signature**\n",
    "\n",
    "在对 tf.function 的调用中指定输入张量，如下所示。此 tf.function 只能在与指定的前面匹配的张量上执行。\n",
    "\n",
    "形状中的 None 用作通配符。这就是 tf.TensroSpec 所说的 \"任意长度的 float32 向量\"。\n",
    "\n",
    "如果你的 tf.function 需要处理不同长度的序列，或者每个批处理不同大小的图像（例如，参见 [Transformer](https://tensorflow.google.cn/tutorials/text/transformer) 和 [Deep Dream tutrorials](https://tensorflow.google.cn/tutorials/generative/deepdream)），则此模式非常重要。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tf.function(input_signature=(\n",
    "    tf.TensorSpec(shape=[None], dtype=tf.float32),\n",
    "    tf.TensorSpec(shape=[None], dtype=tf.float32),)\n",
    ")\n",
    "def power_with_sig(a,b):\n",
    "  print('跟踪 \"power_with_sig\"\\n')\n",
    "  return a**b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 调用 get_concrete_function 将执行跟踪（如果需要），并返回一个具体函数。\n",
    "p = power_with_sig.get_concrete_function()\n",
    "print(type(p))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p(tf.constant([2.0,3.0,4.0]), tf.constant([5.0,4.0,3.0])).numpy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 使用 get_concrete_function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tf.function\n",
    "def power(a,b):\n",
    "  print('跟踪 \"power\"\\n')\n",
    "  return a**b\n",
    "\n",
    "float_power = power.get_concrete_function(\n",
    "  a = tf.TensorSpec(shape=[], dtype=tf.float32),\n",
    "  b = tf.TensorSpec(shape=[], dtype=tf.float32))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "float_power(tf.constant(3.0),tf.constant(3.0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 记住，你还可以传递张量来获取具体函数，在这种情况下，它将返回为这些输入运行的具体函数\n",
    "row = tf.range(10)\n",
    "col = tf.constant([[1],[2],[3]])\n",
    "\n",
    "concrete_power = power.get_concrete_function(a = row, b = col)\n",
    "concrete_power(row, col).numpy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 使用具体函数\n",
    "具体函数只接受张量作为输入："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "float_power(tf.constant(2.0), tf.constant(3.0)).numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "  float_power(2.0,3.0)\n",
    "  assert False\n",
    "except (ValueError, TypeError):\n",
    "  traceback.print_exc(limit=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 也只接受正确类型的张量\n",
    "try:\n",
    "  float_power(tf.constant(1),tf.constant(3))\n",
    "  assert False\n",
    "except tf.errors.InvalidArgumentError:\n",
    "  traceback.print_exc(limit=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 但它将尝试执行，即使输入张量与预期形状不匹配：\n",
    "float_power(tf.constant([1.,2.,3.,4.,5.]),tf.constant(3.)).numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "  float_power(tf.constant([1.,2.,3.]),tf.constant([4., 5.])).numpy()\n",
    "  assert False\n",
    "except tf.errors.InvalidArgumentError:  \n",
    "  traceback.print_exc(limit=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 通过检查具体功能，可以看到其输入和输出\n",
    "print(float_power.structured_input_signature)\n",
    "print(float_power.structured_outputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 签名中的 python 对象\n",
    "正如你在跟踪时看到的，每个 python 对象都会生成一个新的跟踪。具体函数表示一个 tf.Graph，它们不做任何回溯。当使用 python 对象作为参数之一调用 get_concrete_function 时，该对象将绑定到该函数。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cube = power.get_concrete_function(\n",
    "    a = tf.TensorSpec([], dtype=tf.float32),\n",
    "    b = 3.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 此 cube 数据集函数不再具有 b 参数\n",
    "print(cube.structured_input_signature)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cube(tf.constant(10.0)).numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 这与标准 python 类绑定方法的方式非常相似，并且在从方法运行g et_concrete_function 时同样适用\n",
    "class Greeter(object):\n",
    "  def __init__(self, greeting):\n",
    "    self.greeting = greeting\n",
    "\n",
    "  def greet(self, who):\n",
    "    return \" \".join([self.greeting, who])\n",
    "\n",
    "p = Greeter(\"你好\")\n",
    "m = p.greet\n",
    "print(m)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(m(\"TensorFlow!\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 当有t f.function 修饰方法时，也会应用相似的规则\n",
    "class MyModel(tf.Module):\n",
    "  def __init__(self, ins, outs):\n",
    "    initializer = tf.initializers.GlorotNormal()\n",
    "    self.W = tf.Variable(initializer([ins, outs]))\n",
    "    self.B = tf.Variable(tf.zeros([outs], dtype = tf.float32))\n",
    "\n",
    "  @tf.function\n",
    "  def run(self, x):\n",
    "    print('跟踪 \"MyModule\":\\n    x={}\\n'.format(x))\n",
    "    return tf.matmul(x, self.W)+self.B\n",
    "\n",
    "mod = MyModel(ins=5, outs=3)\n",
    "\n",
    "mod.run([[1.0,1.0,1.0, 1.0, 1.0]]).numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 如果你调用它的 .get_concrete_function 方法，self 会自动绑定为第一个参数\n",
    "concrete_run = mod.run.get_concrete_function(x = tf.TensorSpec([None, None]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "concrete_run(tf.constant([[1.0,1.0,1.0, 1.0, 1.0],\n",
    "                          [2.0,2.0,2.0, 2.0, 2.0]])).numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 看看 self 不再是输入签名的一部分\n",
    "print(concrete_run.structured_input_signature)\n",
    "print(concrete_run.structured_outputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 在 SavedModel 中访问具体函数\n",
    "当你保存 SavedModel 时，实际上是在保存 tf.function 的具体函数缓存。\n",
    "\n",
    "因为具体函数是通过跟踪输入生成的，所以至少需要执行一个跟踪来保存 SavedModel。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dropout = Dropout(0.5)\n",
    "\n",
    "_ = dropout(tf.range(10, dtype=tf.float32), tf.constant(True))\n",
    "_ = dropout(tf.random.normal([2, 3]), tf.constant(True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> 注意：在保存时，tf.saved_model 会追溯所有 concrete_function。这是为了确保导出的具体函数捕获导出时环境的变化（例如分发策略范围）。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "export_dir = 'dropout'\n",
    "tf.saved_model.save(dropout, export_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 直接存取\n",
    "当你加载 tf.saved_model 时，方法将还原为多态函数："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reloaded_dropout = tf.saved_model.load(export_dir)\n",
    "\n",
    "print(reloaded_dropout(tf.range(10, dtype=tf.float32), tf.constant(False)).numpy())\n",
    "print(reloaded_dropout(tf.random.normal([2,3]), tf.constant(True)).numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "但是，由于 saved_model 只包含具体函数（而不是 python 源和数据）的缓存，因此它无法处理不匹配的签名："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "  reloaded_dropout(tf.range(12, dtype=tf.float32), tf.constant(True))\n",
    "  assert False\n",
    "except ValueError:\n",
    "  traceback.print_exc(limit=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "从重新加载的模块中，你可以选择一个特定的具体函数，而不是依赖被分配，再次使用get_concrete_function 方法："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cf = reloaded_dropout.__call__.get_concrete_function(\n",
    "    x = tf.TensorSpec([10]), \n",
    "    training = tf.TensorSpec([], tf.bool))\n",
    "\n",
    "result = cf(tf.range(10, dtype=tf.float32), tf.constant(True)).numpy()\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 命名签名：为 C++ 导出\n",
    "使用 SavedModels  的 C++ 用户不用上述直接访问方法，也就是动态调度，而是从 SavedModels 获取和运行具体的函数。\n",
    "\n",
    "它们使用一个更显式的接口，称为 \"导出签名（exported signatures）\"，在这个接口中，你可以确切地指定要导出的具体函数。\n",
    "\n",
    "通过将 signatures 参数传递给 [tf.saved_model.save](https://tensorflow.google.cn/api_docs/python/tf/saved_model/save)，可以指定要导出的具体函数。\n",
    "\n",
    "它需要：\n",
    "- 功能字典。这允许你命名每个函数。\n",
    "- 单一的功能。当一个函数被导出时，它将被命名为 \"serving_default\"，使用 [saved_model.default_serving_SIGNATURE_DEF_KEY](https://tensorflow.google.cn/api_docs/python/tf/saved_model#DEFAULT_SERVING_SIGNATURE_DEF_KEY)。\n",
    "\n",
    "使用 TensorFlow Serving 时需要这些签名。\n",
    "\n",
    "> 注意：这些 \"导出签名（\"exported signatures）\" 被包装以始终返回结果字典。\n",
    "\n",
    "#### 简单的例子"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dropout = Dropout(0.5)\n",
    "\n",
    "cf = dropout.__call__.get_concrete_function(tf.zeros((2,3), dtype=tf.float32), tf.constant(False))\n",
    "\n",
    "import time\n",
    "export_dir = \"./saved/\"+str(time.time())\n",
    "\n",
    "tf.saved_model.save(dropout, export_dir, signatures = cf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "此 saved_model 只包含一个签名，可以从签名字典中按名称恢复:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reloaded = tf.saved_model.load(export_dir)\n",
    "\n",
    "print(reloaded.signatures)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "当使用 \"export signatures\"，这些具体函数始终返回输出字典："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cf = reloaded.signatures[tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY]\n",
    "result = cf(x=tf.random.normal([2,3]), training=tf.constant(True))\n",
    "\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在上面的示例中，由签名自动生成的输出名称很通用。可以使用结构化输出方法检查输出名称：\n",
    "\n",
    "你可以使用 .structured_outputs 方法检查预期的输出张量名称："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cf.structured_outputs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "通常你不想自己设置输出名称。\n",
    "\n",
    "#### 示例：设置输出名称\n",
    "\n",
    "要控制输出的名称，修改 tf.function 以返回将名称映射到输出张量的字典："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tf.function\n",
    "def named_result(x, training=True):\n",
    "  return {'dropout': dropout(x, training)}\n",
    "\n",
    "dropout.named_result = named_result\n",
    "\n",
    "cf = dropout.named_result.get_concrete_function(tf.zeros((2,3), dtype=tf.float32),\n",
    "                                                tf.constant(False))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 示例：设置签名名称\n",
    "若要设置签名的名称，传递一个包含具体函数的字典。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "export_dir = \"./saved/\"+str(time.time())\n",
    "tf.saved_model.save(dropout, export_dir, signatures = {'simple':cf})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reloaded = tf.saved_model.load(export_dir)\n",
    "cf = reloaded.signatures['simple']\n",
    "result = cf(x=tf.random.normal([2,3]), training=tf.constant(True))\n",
    "\n",
    "print({key:value.numpy() for key,value in result.items()})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "若要指定多个签名，将 (name, concrete_function) 对的字典传递给 [save_model.save](https://tensorflow.google.cn/api_docs/python/tf/saved_model/save)："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vector = dropout.__call__.get_concrete_function(tf.TensorSpec((2,3), dtype=tf.float32), tf.constant(False))\n",
    "matrix = dropout.__call__.get_concrete_function(tf.TensorSpec((2,3), dtype=tf.float32), tf.constant(False))\n",
    "cube = dropout.__call__.get_concrete_function(tf.TensorSpec((2,3), dtype=tf.float32), tf.constant(False))\n",
    "\n",
    "export_dir = \"./saved/\"+str(time.time())\n",
    "\n",
    "tf.saved_model.save(dropout, export_dir, \n",
    "                    signatures = {\n",
    "                        \"vector\": vector,\n",
    "                        \"matrix\": matrix,\n",
    "                        \"cube\": cube\n",
    "                    })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 现在加载模型和检查签名列表\n",
    "reloaded = tf.saved_model.load(export_dir)\n",
    "print('{}'.format(reloaded.signatures).replace(\"{\",\"{\\n    \").replace(\">, \", \">,\\n    \"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 在 tensorflow2.3 中的改变\n",
    "以下改变在 tensorflow 最新版可用，而且在 tensorflow2.3 中可用。\n",
    "\n",
    "### 检查具体函数签名\n",
    "打印 ConcreteFunction 将显示其输入参数（带类型）及其输出类型的摘要："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(float_power)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "对于多态函数，pretty_printed_concrete_signatures() 方法可用于显示目前已跟踪的所有签名的列表："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(power.pretty_printed_concrete_signatures())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 绑定参数\n",
    "如上所述，在创建具体函数时，设置为非张量值的任何参数都将绑定到这些值。在 TensorFlow2.3 之前，这些绑定参数只是从具体函数的签名中移除。从 TensorFlow2.3 开始，它们仍然存在于签名中，但其默认值设置为绑定值。例如，当我们创建以下具体函数时，参数 b 的默认值设置为 2："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "square = power.get_concrete_function(a=tf.TensorSpec(None, tf.float32), b=2)\n",
    "print(square)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "调用 square 时，可以省略 b，也可以将其显式设置为跟踪期间使用的值："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(square(tf.constant(5.0)))\n",
    "print(square(tf.constant(5.0), b=2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "将绑定参数设置为默认值以外的任何值都是错误的："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "  print(square(tf.constant(10.0), b=4.0).numpy())  # 错误:显式值!=绑定值\n",
    "  assert False\n",
    "except TypeError:\n",
    "  traceback.print_exc(limit=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### CompositeTensor  参数\n",
    "从 tensorflow2.3 开始，你可以定义接受 CompositeTensors 的具体函数（例如 [tf.raggedSensor](https://tensorflow.google.cn/api_docs/python/tf/RaggedTensor)、tf.Dataset迭代器和tf.SparseTensor）。例如："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ragged_power = power.get_concrete_function(\n",
    "    a = tf.RaggedTensorSpec([None, None], dtype=tf.float32),\n",
    "    b = tf.TensorSpec([], dtype=tf.float32)\n",
    ")\n",
    "print(ragged_power(tf.ragged.constant([[1.0, 2.0], [3.0]]),\n",
    "                   tf.constant(2.0)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 内嵌参数\n",
    "从 tensorflow2.3 开始，你可以定义具体的函数，这些函数接受张量（或复合张量）的嵌套 dict、list、tuple、namedtuples 和 attrs，例如："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tf.function\n",
    "def sum_features(features):\n",
    "  assert isinstance(features, dict)\n",
    "  return sum(t for t in features.values())\n",
    "\n",
    "features = {\n",
    "    'a': tf.constant([[1.0], [2.0], [8.0]]),\n",
    "    'b': tf.constant([[5.0], [6.0], [0.0]]),\n",
    "    'c': tf.ragged.constant([[10.0], [20.0, 30.0, 40.0], [50.0]])\n",
    "}\n",
    "print(sum_features(features))\n",
    "\n",
    "concrete_sum_features = sum_features.get_concrete_function(features)\n",
    "print(concrete_sum_features(features))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "传递给具体函数的值必须与跟踪具体函数时使用的嵌套结构相同。如果传入其他结构，则会引发T ypeError，其中包含有关已接受的结构和已使用的结构的信息："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "  concrete_sum_features({'a': tf.constant([[1.0]])})\n",
    "  assert False\n",
    "except TypeError:\n",
    "  traceback.print_exc(limit=1)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6rc1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
